{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker V3 JumpStart Training Example\n",
    "\n",
    "This notebook demonstrates how to use SageMaker V3 ModelTrainer with JumpStart models for easy model training and fine-tuning."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "Note: Ensure you have sagemaker-train and ipywidgets installed in your environment. The ipywidgets package is required to monitor training job progress in Jupyter notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import json\n",
    "import uuid\n",
    "\n",
    "from sagemaker.train.model_trainer import ModelTrainer\n",
    "from sagemaker.core.jumpstart import JumpStartConfig\n",
    "from sagemaker.core.helper.session_helper import Session, get_execution_role"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Setup Session and Configuration\n",
    "\n",
    "Initialize the SageMaker session and define our training configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize SageMaker session\n",
    "sagemaker_session = Session()\n",
    "role = get_execution_role()\n",
    "\n",
    "# Configuration\n",
    "JOB_NAME_PREFIX = \"js-v3-training-example\"\n",
    "\n",
    "# Generate unique identifier\n",
    "unique_id = str(uuid.uuid4())[:8]\n",
    "base_job_name = f\"{JOB_NAME_PREFIX}-{unique_id}\"\n",
    "\n",
    "print(f\"Base job name: {base_job_name}\")\n",
    "print(f\"SageMaker execution role: {role}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Train HuggingFace BERT Model\n",
    "\n",
    "Train a HuggingFace BERT model for text classification using JumpStart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure JumpStart for HuggingFace BERT\n",
    "bert_jumpstart_config = JumpStartConfig(\n",
    "    model_id=\"huggingface-spc-bert-base-cased\",\n",
    "    accept_eula=False  # This model doesn't require EULA acceptance\n",
    ")\n",
    "\n",
    "# Create ModelTrainer from JumpStart config\n",
    "bert_trainer = ModelTrainer.from_jumpstart_config(\n",
    "    jumpstart_config=bert_jumpstart_config,\n",
    "    base_job_name=f\"{base_job_name}-bert\",\n",
    "    hyperparameters={\n",
    "        \"epochs\": 1,  # Set to 1 for quick demonstration\n",
    "        \"learning_rate\": 5e-5,\n",
    "        \"train_batch_size\": 32\n",
    "    },\n",
    "    sagemaker_session=sagemaker_session\n",
    ")\n",
    "\n",
    "print(\"BERT ModelTrainer created successfully from JumpStart config!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start BERT training\n",
    "print(\"Starting BERT training job...\")\n",
    "print(\"Note: This will use the default JumpStart dataset and may take 10-15 minutes.\")\n",
    "\n",
    "bert_trainer.train()\n",
    "print(f\"BERT training job completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Train XGBoost Classification Model\n",
    "\n",
    "Train an XGBoost model for classification tasks using JumpStart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure JumpStart for XGBoost\n",
    "xgboost_jumpstart_config = JumpStartConfig(\n",
    "    model_id=\"xgboost-classification-model\"\n",
    ")\n",
    "\n",
    "# Create ModelTrainer from JumpStart config\n",
    "xgboost_trainer = ModelTrainer.from_jumpstart_config(\n",
    "    jumpstart_config=xgboost_jumpstart_config,\n",
    "    base_job_name=f\"{base_job_name}-xgboost\",\n",
    "    hyperparameters={\n",
    "        \"num_round\": 10,  # Reduced for quick demonstration\n",
    "        \"max_depth\": 5,\n",
    "        \"eta\": 0.2,\n",
    "        \"objective\": \"binary:logistic\"\n",
    "    },\n",
    "    sagemaker_session=sagemaker_session\n",
    ")\n",
    "\n",
    "print(\"XGBoost ModelTrainer created successfully from JumpStart config!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start XGBoost training\n",
    "print(\"Starting XGBoost training job...\")\n",
    "print(\"Note: This will use the default JumpStart dataset and should complete in 5-10 minutes.\")\n",
    "\n",
    "xgboost_trainer.train()\n",
    "print(f\"XGBoost training job completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Train CatBoost Regression Model\n",
    "\n",
    "Train a CatBoost model for regression tasks using JumpStart."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure JumpStart for CatBoost\n",
    "catboost_jumpstart_config = JumpStartConfig(\n",
    "    model_id=\"catboost-regression-model\"\n",
    ")\n",
    "\n",
    "# Create ModelTrainer from JumpStart config\n",
    "catboost_trainer = ModelTrainer.from_jumpstart_config(\n",
    "    jumpstart_config=catboost_jumpstart_config,\n",
    "    base_job_name=f\"{base_job_name}-catboost\",\n",
    "    hyperparameters={\n",
    "        \"iterations\": 50,  # Reduced for quick demonstration\n",
    "        \"learning_rate\": 0.1,\n",
    "        \"depth\": 6,\n",
    "        \"loss_function\": \"RMSE\"\n",
    "    },\n",
    "    sagemaker_session=sagemaker_session\n",
    ")\n",
    "\n",
    "print(\"CatBoost ModelTrainer created successfully from JumpStart config!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start CatBoost training\n",
    "print(\"Starting CatBoost training job...\")\n",
    "print(\"Note: This will use the default JumpStart dataset and should complete in 5-10 minutes.\")\n",
    "\n",
    "catboost_trainer.train()\n",
    "print(f\"CatBoost training job completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Review Training Results\n",
    "\n",
    "Check the status and results of our training jobs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display training job information\n",
    "training_jobs = [\n",
    "    (\"BERT\", bert_trainer),\n",
    "    (\"XGBoost\", xgboost_trainer),\n",
    "    (\"CatBoost\", catboost_trainer)\n",
    "]\n",
    "\n",
    "print(\"Training Job Summary:\")\n",
    "print(\"=\" * 50)\n",
    "\n",
    "for model_name, trainer in training_jobs:\n",
    "    job_name = trainer._latest_training_job.training_job_name\n",
    "    model_artifacts = trainer._latest_training_job.model_artifacts\n",
    "    \n",
    "    print(f\"\\n{model_name} Model:\")\n",
    "    print(f\"  Job Name: {job_name}\")\n",
    "    print(f\"  Model Artifacts: {model_artifacts}\")\n",
    "    print(f\"  Status: Completed\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Access Training Metrics (Optional)\n",
    "\n",
    "View training metrics and logs from CloudWatch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example: Access training job details\n",
    "print(\"Training Job Details:\")\n",
    "print(\"\\nTo view detailed training metrics and logs:\")\n",
    "print(\"1. Go to the SageMaker Console\")\n",
    "print(\"2. Navigate to 'Training' > 'Training jobs'\")\n",
    "print(\"3. Search for jobs with prefix:\", base_job_name)\n",
    "print(\"4. Click on any job to view metrics, logs, and model artifacts\")\n",
    "\n",
    "# You can also access logs programmatically\n",
    "print(\"\\nProgrammatic access to logs:\")\n",
    "for model_name, trainer in training_jobs:\n",
    "    print(f\"{model_name}: trainer.latest_training_job.describe()\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "1. Creating ModelTrainer instances from JumpStart configurations\n",
    "2. Training multiple model types (BERT, XGBoost, CatBoost) with custom hyperparameters\n",
    "3. Using JumpStart's built-in datasets and training scripts\n",
    "4. Monitoring training job progress and results\n",
    "\n",
    "## Benefits of JumpStart Training:\n",
    "- **Pre-configured models**: No need to write training scripts or handle data preprocessing\n",
    "- **Best practices**: Optimized hyperparameters and training configurations\n",
    "- **Multiple frameworks**: Support for HuggingFace, XGBoost, CatBoost, and more\n",
    "- **Easy customization**: Override hyperparameters while keeping proven defaults\n",
    "- **Built-in datasets**: Start training immediately with curated datasets\n",
    "\n",
    "## Next Steps:\n",
    "- Deploy trained models using SageMaker V3 ModelBuilder\n",
    "- Fine-tune models with your own datasets\n",
    "- Experiment with different hyperparameters\n",
    "- Set up automated training pipelines\n",
    "\n",
    "JumpStart training with V3 ModelTrainer makes it incredibly easy to get started with machine learning while maintaining the flexibility to customize as needed!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
